# 导入操作系统相关的功能模块，用于文件路径操作等
import os

# 导入 OpenCV 库，用于图像处理
import cv2
# 导入 NumPy 库，用于数值计算
import numpy as np
# 导入 PyTorch 深度学习框架
import torch
# 从 PIL 库中导入 Image 类，用于图像的读取和处理
from PIL import Image

# 从自定义的 dataset 模块中导入 image_square_pad 函数，用于将图像填充为正方形
from dataset import image_square_pad
# 从自定义的 model 模块中导入 SwitchModel 类，用于构建模型
from model import SwitchModel


def create_model(device):
    # 检查是否有可用的 CUDA 设备，如果有则使用 GPU，否则使用 CPU
    # 定义分类标签列表
    labels = ["OFF", "ON"]
    # 定义分类的类别数量
    num_classes = 2
    # 定义模型使用的骨干网络类型
    backbone = "shufflenet"
    # 定义模型权重文件的路径
    model_path = "output/model_shufflenet.pth"
    
    # 创建 SwitchModel 实例，不使用预训练权重
    model = SwitchModel(backbone=backbone, num_classes=num_classes, pretrained=False)
    # 加载模型的权重文件，指定加载到 CPU 上
    model.load_state_dict(torch.load(model_path, map_location=torch.device("cpu"), weights_only=True))
    # 将模型设置为评估模式
    model.eval()
    # 将模型移动到指定的设备上
    model.to(device)

    model.labels = labels

    # raw run
    raw_image = np.zeros((224, 224, 3), dtype=np.uint8)
    result = run_model(model, device, raw_image)

    return model

def run_model(model, device, image):
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    # 调用 image_square_pad 函数将图像填充为正方形
    image, _ = image_square_pad(image)
    # 将图像调整为 224x224 的大小
    image = cv2.resize(image, (224, 224)) 
    
    # 将 NumPy 数组转换为 PyTorch 张量，并调整维度顺序
    image = torch.from_numpy(image).permute(2, 0, 1).unsqueeze(0).float()
    # 对图像进行归一化处理
    image = image / 255.0 * 2 - 1
    # 将图像张量移动到指定的设备上
    image = image.to(device)
    
    # 禁用梯度计算，减少内存消耗
    with torch.no_grad():
        # 前向传播，通过模型得到输出
        output = model(image)
        # 获取预测结果的索引
        _, predicted = torch.max(output, 1)
    
    result = model.labels[predicted.item()]

    return result


    